import time
import datetime
from typing import List

import joblib
import pandas as pd

from base_funcs import *
from base_model import init_base_model
from dynamic_model import init_dynamic_model


if __name__ == "__main__":
    calc_begin_time = time.perf_counter()
    result = {
        "h_m": {},
        "c_e": {},
        "h_e": {}
    }
    # h_m
    for idx, h_m in enumerate([0.1, 0.15, 0.2, 0.25, 0.3]):
        params = {"_T": 4, "_b": 20, "_c_m": 2, "_h_m": h_m, "_c_e": 10, "_h_e": 1}
        res_base_model = init_base_model(**params)
        res_dynamic_model = init_dynamic_model(**params)
        res_relative_error = relative_error(res_base_model, res_dynamic_model)
        result["h_m"].update({f"{idx + 1}": {"base_model": res_base_model,
                                             "dynamic_model": res_dynamic_model,
                                             "relative_error": res_relative_error,
                                             "params": params
                                             }})
    # c_e
    for idx, c_e in enumerate([6, 8, 10, 12, 14]):
        params = {"_T": 4, "_b": 20, "_c_m": 2, "_h_m": 0.2, "_c_e": c_e, "_h_e": 1}
        res_base_model = init_base_model(**params)
        res_dynamic_model = init_dynamic_model(**params)
        res_relative_error = relative_error(res_base_model, res_dynamic_model)
        result["c_e"].update({f"{idx + 1}": {"base_model": res_base_model,
                                             "dynamic_model": res_dynamic_model,
                                             "relative_error": res_relative_error,
                                             "params": params
                                             }})

    # h_e
    for idx, h_e in enumerate([0.6, 0.8, 1, 1.2, 1.4]):
        params = {"_T": 4, "_b": 20, "_c_m": 2, "_h_m": 0.2, "_c_e": 10, "_h_e": h_e}
        res_base_model = init_base_model(**params)
        res_dynamic_model = init_dynamic_model(**params)
        res_relative_error = relative_error(res_base_model, res_dynamic_model)
        result["h_e"].update({f"{idx + 1}": {"base_model": res_base_model,
                                             "dynamic_model": res_dynamic_model,
                                             "relative_error": res_relative_error,
                                             "params": params
                                             }})

    calc_end_time = time.perf_counter()
    print(f"计算完成, 耗时 {calc_end_time - calc_begin_time: .2f} 秒")
    # output

    writer = pd.ExcelWriter(f"{datetime.datetime.now().strftime('%Y-%m-%d %H-%M-%S')}.xlsx")
    data_h_m = pd.DataFrame(result["h_m"])
    data_c_e = pd.DataFrame(result["c_e"])
    data_h_e = pd.DataFrame(result["h_e"])
    data_h_m.to_excel(writer, sheet_name="h_m")
    data_c_e.to_excel(writer, sheet_name="c_e")
    data_h_e.to_excel(writer, sheet_name="h_e")
    writer.close()
    joblib.dump(result, f"{datetime.datetime.now().strftime('%Y-%m-%d %H-%M-%S')}.pkl")
